package net.aiion.weave.impl.interception.transformation;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;

import net.aiion.weave.asm.Label;
import net.aiion.weave.asm.Opcodes;
import net.aiion.weave.asm.Type;
import net.aiion.weave.spi.bytecode.Bytecode;
import net.aiion.weave.spi.bytecode.BytecodeContext;
import net.aiion.weave.spi.bytecode.BytecodeException;
import net.aiion.weave.spi.bytecode.BytecodeGenerator;
import net.aiion.weave.spi.bytecode.events.AnnotationEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.AnnotationStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.AnnotationValueBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.BytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.util.BytecodeTransformation;
import net.aiion.weave.spi.interception.InterceptionContext;
import net.aiion.weave.spi.interception.Interceptor;
import net.aiion.weave.spi.interception.annotations.Intercepts;

// TODO Add inheritance support (also respect annotations of implemented interfaces).
// TODO Add real support for multiple @Intercepts annotated methods. Currently, a call to InterceptionContext.forwardInterception() skips all other possibly available @Intercepts methods in the same class. 
public class InterceptorTransformation extends BytecodeTransformation {
    private static final String INTERCEPTION_CONTEXT_INTERNAL_NAME = Type.getInternalName(InterceptionContext.class);

    private static final String INTERCEPTOR_INTERCEPT_INVOCATION_DESC = "(L" + INTERCEPTION_CONTEXT_INTERNAL_NAME
        + ";)Ljava/lang/Object;";

    private String className;

    private String methodName;
    private String methodDescriptor;
    private Type interceptedAnnotation;

    private boolean inInterceptsAnnotation;

    private Set<Type> interceptedAnnotations;

    private Collection<BytecodeEvent> interceptInvocationBytecodeEvents;
    private BytecodeGenerator interceptInvocationBytecodeGenerator;

    public InterceptorTransformation() {
    }

    @Override
    public boolean testRetain(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && !aBytecodeEvent.asClassStart().containsInterface(Interceptor.class);
    }

    @Override
    public boolean testProcess(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.ANNOTATION_START
            && aBytecodeEvent.asAnnotationStart().equalsType(Intercepts.class);
    }

    @Override
    public boolean testSkip(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && (aBytecodeEvent.asClassStart().isInterface() || aBytecodeEvent.asClassStart().containsInterface(
                Interceptor.class));
    }

    @Override
    public void handleBytecodeEvent(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        switch (aBytecodeEvent.getType()) {
        case BytecodeEvent.CLASS_START:
            handleClassStartBytecode(aBytecodeContext, aBytecodeEvent.asClassStart());
            break;
        case BytecodeEvent.METHOD_START:
            handleMethodStartBytecode(aBytecodeContext, aBytecodeEvent.asMethodStart());
            break;
        case BytecodeEvent.ANNOTATION_START:
            handleAnnotationStartBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationStart());
            break;
        case BytecodeEvent.ANNOTATION_VALUE:
            handleAnnotationValueBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationValue());
            break;
        case BytecodeEvent.ANNOTATION_END:
            handleAnnotationEndBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationEnd());
            break;
        case BytecodeEvent.METHOD_END:
            handleMethodEndBytecode(aBytecodeContext, aBytecodeEvent.asMethodEnd());
            break;
        case BytecodeEvent.CLASS_END:
            handleClassEndBytecode(aBytecodeContext, aBytecodeEvent.asClassEnd());
            break;
        default:
            break;
        }

        aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
    }

    private void handleClassStartBytecode(final BytecodeContext aBytecodeContext,
        final ClassStartBytecodeEvent aClassStartBytecodeEvent) {
        className = aClassStartBytecodeEvent.getName();
        aClassStartBytecodeEvent.addInterface(Interceptor.class);

        prepareInterceptInvocation();
    }

    private void handleMethodStartBytecode(final BytecodeContext aBytecodeContext,
        final MethodStartBytecodeEvent aMethodStartBytecodeEvent) {
        methodName = aMethodStartBytecodeEvent.getName();
        methodDescriptor = aMethodStartBytecodeEvent.getDescriptor();
    }

    private void handleAnnotationStartBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationStartBytecodeEvent aAnnotationStartBytecodeEvent) {
        if (methodName != null && aAnnotationStartBytecodeEvent.equalsType(Intercepts.class)) {
            inInterceptsAnnotation = true;

            testInterceptorMethod();
        }
    }

    public void testInterceptorMethod() {
        if (!methodDescriptor.equals(INTERCEPTOR_INTERCEPT_INVOCATION_DESC)) {
            throw new BytecodeException("A @Intercepts annotated method must accept exactly one argument of type "
                + InterceptionContext.class.getName() + " and must not return java.lang.Object.");
        }
    }

    private void handleAnnotationValueBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationValueBytecodeEvent aAnnotationValueBytecodeEvent) {
        if (inInterceptsAnnotation && "value".equals(aAnnotationValueBytecodeEvent.getName())) {
            interceptedAnnotation = (Type) aAnnotationValueBytecodeEvent.getValue();
        }
    }

    private void handleAnnotationEndBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationEndBytecodeEvent aAnnotationEndBytecodeEvent) {
        inInterceptsAnnotation = false;
    }

    private void handleMethodEndBytecode(final BytecodeContext aBytecodeContext,
        final MethodEndBytecodeEvent aMethodEndBytecodeEvent) {
        if (interceptedAnnotation != null) {
            extendInterceptInvocation();
        }

        methodName = null;
        interceptedAnnotation = null;
    }

    private void handleClassEndBytecode(final BytecodeContext aBytecodeContext,
        final ClassEndBytecodeEvent aClassEndBytecodeEvent) {
        finishInterceptInvocation();

        aBytecodeContext.forwardBytecodeEvents(interceptInvocationBytecodeEvents);
    }

    private void prepareInterceptInvocation() {
        interceptInvocationBytecodeEvents = new ArrayList<>();
        interceptInvocationBytecodeGenerator = Bytecode.createGenerator(interceptInvocationBytecodeEvents);

        interceptInvocationBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "interceptInvocation", Object.class,
            InterceptionContext.class);
        interceptInvocationBytecodeGenerator.startCode();

        interceptInvocationBytecodeGenerator.loadReference(1);
        interceptInvocationBytecodeGenerator.invokeInterface(InterceptionContext.class, "getInterceptedMethod",
            Method.class);
        interceptInvocationBytecodeGenerator.storeReference(2);

        interceptedAnnotations = new HashSet<>();
    }

    private void extendInterceptInvocation() {
        if (interceptedAnnotations.contains(interceptedAnnotation)) {
            throw new BytecodeException(
                "There must be only one @Intercepts annotated method for intercepted annotation type "
                    + interceptedAnnotation.getClassName() + " in " + className);
        }

        if (!interceptedAnnotations.isEmpty()) {
            interceptInvocationBytecodeGenerator.markFrameAppend(1, new Object[] { "java/lang/reflect/Method" }, 0,
                null);
        }

        final Label nextAnnotationLabel = new Label();

        interceptInvocationBytecodeGenerator.loadReference(2);
        interceptInvocationBytecodeGenerator.pushConstant(interceptedAnnotation);
        interceptInvocationBytecodeGenerator.invokeVirtual(Method.class, "isAnnotationPresent", boolean.class,
            Class.class);
        interceptInvocationBytecodeGenerator.jumpEquals(nextAnnotationLabel);

        interceptInvocationBytecodeGenerator.loadReference(0);
        interceptInvocationBytecodeGenerator.loadReference(1);
        interceptInvocationBytecodeGenerator.invokeVirtual(className, methodName, methodDescriptor);
        interceptInvocationBytecodeGenerator.returnReference(Type.getReturnType(methodDescriptor));
        interceptInvocationBytecodeGenerator.markLabel(nextAnnotationLabel);

        interceptedAnnotations.add(interceptedAnnotation);
    }

    private void finishInterceptInvocation() {
        interceptInvocationBytecodeGenerator.markFrameSame();

        interceptInvocationBytecodeGenerator.loadReference(1);
        interceptInvocationBytecodeGenerator.invokeInterface(InterceptionContext.class, "forwardInterception",
            Object.class);
        interceptInvocationBytecodeGenerator.returnReference();

        interceptInvocationBytecodeGenerator.endCode(2, 3);
        interceptInvocationBytecodeGenerator.endMethod();

        interceptedAnnotations.clear();
    }
}
